# file: utils/file_loader.py

import os

from langchain_community.document_loaders import PyPDFLoader, Docx2txtLoader, TextLoader


def load_documents(data_dir="data"):
    """
    加载多种格式的金融文档

    Args:
        data_dir (str): 文档目录路径，默认为"data"

    Returns:
        list: 文档对象列表
    """
    documents = []

    # 检查数据目录是否存在
    if not os.path.exists(data_dir):
        print(f"警告: 目录 {data_dir} 不存在，将创建空目录")
        os.makedirs(data_dir)
        return documents

    # 支持的文件扩展名
    loaders = {
        '.pdf': PyPDFLoader,
        '.docx': Docx2txtLoader,
        '.txt': TextLoader
    }

    # 遍历数据目录中的所有文件
    for file_name in os.listdir(data_dir):
        file_path = os.path.join(data_dir, file_name)

        # 只处理文件（跳过目录）
        if os.path.isfile(file_path):
            file_extension = os.path.splitext(file_name)[1].lower()

            # 根据文件扩展名选择合适的加载器
            if file_extension in loaders:
                try:
                    loader = loaders[file_extension](file_path)
                    if file_extension == '.txt':
                        # TextLoader需要指定编码
                        loader = TextLoader(file_path, encoding='utf-8')
                        docs = loader.load()
                    else:
                        docs = loader.load()

                    # 为每个文档添加元数据
                    for doc in docs:
                        doc.metadata['source'] = file_name

                    documents.extend(docs)
                    print(f"成功加载文件: {file_name}")

                except Exception as e:
                    print(f"加载文件 {file_name} 时出错: {str(e)}")
            else:
                print(f"跳过不支持的文件类型: {file_name}")

    print(f"总共加载了 {len(documents)} 个文档片段")
    return documents
